import Foundation
import Combine

class OpenAIService {
    static let shared = OpenAIService()
    
    private let baseURL = "https://api.openai.com/v1/chat/completions"
    
    private init() {
        // 确保API密钥已保存
        let apiKey = "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"
        KeychainService.saveAPIKey(apiKey, service: "openai")
    }
    
    private func getAPIKey() -> String? {
        return KeychainService.getAPIKey(service: "openai")
    }
    
    // 发送消息并获取响应
    func sendMessage(messages: [ChatMessage], systemPrompt: String, completion: @escaping (Result<String, Error>) -> Void) {
        // 获取API密钥
        guard let apiKey = getAPIKey() else {
            completion(.failure(NSError(domain: "OpenAIService", code: 0, userInfo: [NSLocalizedDescriptionKey: "未找到API密钥"])))
            return
        }
        
        // 准备请求
        guard let url = URL(string: baseURL) else {
            completion(.failure(NSError(domain: "OpenAIService", code: 0, userInfo: [NSLocalizedDescriptionKey: "无效的URL"])))
            return
        }
        
        var request = URLRequest(url: url)
        request.httpMethod = "POST"
        request.addValue("application/json", forHTTPHeaderField: "Content-Type")
        request.addValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization")
        
        // 准备消息数组，包括系统提示
        var apiMessages: [[String: Any]] = []
        
        // 添加系统消息
        if !systemPrompt.isEmpty {
            apiMessages.append([
                "role": "system",
                "content": systemPrompt
            ])
        }
        
        // 添加对话消息
        for message in messages {
            if message.role != .system { // 系统消息已经在上面添加了
                apiMessages.append([
                    "role": message.role.rawValue,
                    "content": message.content
                ])
            }
        }
        
        // 准备请求体
        let requestBody: [String: Any] = [
            "model": "gpt-4o",
            "messages": apiMessages,
            "temperature": 0.7
        ]
        
        do {
            request.httpBody = try JSONSerialization.data(withJSONObject: requestBody)
        } catch {
            completion(.failure(error))
            return
        }
        
        // 创建并启动任务
        let task = URLSession.shared.dataTask(with: request) { data, response, error in
            if let error = error {
                DispatchQueue.main.async {
                    completion(.failure(error))
                }
                return
            }
            
            guard let data = data else {
                DispatchQueue.main.async {
                    completion(.failure(NSError(domain: "OpenAIService", code: 0, userInfo: [NSLocalizedDescriptionKey: "没有数据返回"])))
                }
                return
            }
            
            do {
                let json = try JSONSerialization.jsonObject(with: data, options: []) as? [String: Any]
                if let choices = json?["choices"] as? [[String: Any]],
                   let firstChoice = choices.first,
                   let message = firstChoice["message"] as? [String: Any],
                   let content = message["content"] as? String {
                    DispatchQueue.main.async {
                        completion(.success(content))
                    }
                } else {
                    DispatchQueue.main.async {
                        completion(.failure(NSError(domain: "OpenAIService", code: 0, userInfo: [NSLocalizedDescriptionKey: "无法解析响应"])))
                    }
                }
            } catch {
                DispatchQueue.main.async {
                    completion(.failure(error))
                }
            }
        }
        
        task.resume()
    }
    //流式输出在这里
    
    // 发送消息并获取流式响应
    func sendStreamMessage(messages: [ChatMessage], systemPrompt: String, onReceive: @escaping (String) -> Void, onComplete: @escaping (Error?) -> Void) {
        // 获取API密钥
        guard let apiKey = getAPIKey() else {
            onComplete(NSError(domain: "OpenAIService", code: 0, userInfo: [NSLocalizedDescriptionKey: "未找到API密钥"]))
            return
        }
        
        // 准备请求
        guard let url = URL(string: baseURL) else {
            onComplete(NSError(domain: "OpenAIService", code: 0, userInfo: [NSLocalizedDescriptionKey: "无效的URL"]))
            return
        }
        
        var request = URLRequest(url: url)
        request.httpMethod = "POST"
        request.addValue("application/json", forHTTPHeaderField: "Content-Type")
        request.addValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization")
        
        // 准备消息数组，包括系统提示
        var apiMessages: [[String: Any]] = []
        
        // 添加系统消息
        if !systemPrompt.isEmpty {
            apiMessages.append([
                "role": "system",
                "content": systemPrompt
            ])
        }
        
        // 添加对话消息
        for message in messages {
            if message.role != .system { // 系统消息已经在上面添加了
                apiMessages.append([
                    "role": message.role.rawValue,
                    "content": message.content
                ])
            }
        }
        
        // 准备请求体
        let requestBody: [String: Any] = [
            "model": "gpt-4o",
            "messages": apiMessages,
            "stream": true,
            "temperature": 0.7
        ]
        
        do {
            request.httpBody = try JSONSerialization.data(withJSONObject: requestBody)
        } catch {
            onComplete(error)
            return
        }
        
        // 创建会话
        let session = URLSession(configuration: .default, delegate: nil, delegateQueue: .main)
        
        // 创建数据任务
        let task = session.dataTask(with: request) { data, response, error in
            if let error = error {
                onComplete(error)
                return
            }
            
            guard let httpResponse = response as? HTTPURLResponse else {
                onComplete(NSError(domain: "OpenAIService", code: 0, userInfo: [NSLocalizedDescriptionKey: "无效的HTTP响应"]))
                return
            }
            
            guard (200...299).contains(httpResponse.statusCode) else {
                onComplete(NSError(domain: "OpenAIService", code: httpResponse.statusCode, userInfo: [NSLocalizedDescriptionKey: "HTTP错误: \(httpResponse.statusCode)"]))
                return
            }
            
            guard let data = data else {
                onComplete(NSError(domain: "OpenAIService", code: 0, userInfo: [NSLocalizedDescriptionKey: "没有数据返回"]))
                return
            }
            
            // 处理流式数据
            if let text = String(data: data, encoding: .utf8) {
                let lines = text.components(separatedBy: "\n")
                
                for line in lines {
                    if line.hasPrefix("data: ") {
                        let jsonString = line.dropFirst(6) // 去掉 "data: "
                        
                        if jsonString == "[DONE]" {
                            onComplete(nil)
                            continue
                        }
                        
                        do {
                            if let data = String(jsonString).data(using: .utf8),
                               let json = try JSONSerialization.jsonObject(with: data) as? [String: Any],
                               let choices = json["choices"] as? [[String: Any]],
                               let choice = choices.first,
                               let delta = choice["delta"] as? [String: Any],
                               let content = delta["content"] as? String {
                                onReceive(content)
                            }
                        } catch {
                            print("解析JSON错误: \(error)")
                        }
                    }
                }
            }
            
            onComplete(nil)
        }
        
        task.resume()
    }
    
    // 使用真正的流式处理
    func sendRealTimeStreamMessage(messages: [ChatMessage], systemPrompt: String, modelType: AIModelType = .gpt4o, onReceive: @escaping (String) -> Void, onComplete: @escaping (Error?) -> Void) {
        // 获取API密钥
        guard let apiKey = getAPIKey() else {
            onComplete(NSError(domain: "OpenAIService", code: 0, userInfo: [NSLocalizedDescriptionKey: "未找到API密钥"]))
            return
        }
        
        // 准备请求
        guard let url = URL(string: baseURL) else {
            onComplete(NSError(domain: "OpenAIService", code: 0, userInfo: [NSLocalizedDescriptionKey: "无效的URL"]))
            return
        }
        
        var request = URLRequest(url: url)
        request.httpMethod = "POST"
        request.addValue("application/json", forHTTPHeaderField: "Content-Type")
        request.addValue("Bearer \(apiKey)", forHTTPHeaderField: "Authorization")
        
        // 准备消息数组，包括系统提示
        var apiMessages: [[String: Any]] = []
        
        // 添加系统消息
        if !systemPrompt.isEmpty {
            apiMessages.append([
                "role": "system",
                "content": systemPrompt
            ])
        }
        
        // 添加对话消息
        for message in messages {
            if message.role != .system { // 系统消息已经在上面添加了
                apiMessages.append([
                    "role": message.role.rawValue,
                    "content": message.content
                ])
            }
        }
        
        // 准备请求体，使用传入的模型类型
        let requestBody: [String: Any] = [
            "model": modelType.rawValue,
            "messages": apiMessages,
            "stream": true,
            "temperature": 0.7
        ]
        
        do {
            request.httpBody = try JSONSerialization.data(withJSONObject: requestBody)
        } catch {
            onComplete(error)
            return
        }
        
        // 创建会话和任务
        let streamDelegate = StreamDelegate(onReceive: onReceive, onComplete: onComplete)
        let session = URLSession(configuration: .default, delegate: streamDelegate, delegateQueue: .main)
        let task = session.dataTask(with: request)
        
        // 保存任务引用，防止被释放
        streamDelegate.task = task
        
        task.resume()
    }
}

// 自定义URLSessionDataDelegate来处理流式数据
class StreamDelegate: NSObject, URLSessionDataDelegate {
    private let onReceive: (String) -> Void
    private let onComplete: (Error?) -> Void
    private var buffer = Data()
    var task: URLSessionDataTask?
    
    init(onReceive: @escaping (String) -> Void, onComplete: @escaping (Error?) -> Void) {
        self.onReceive = onReceive
        self.onComplete = onComplete
        super.init()
    }
    
    func urlSession(_ session: URLSession, dataTask: URLSessionDataTask, didReceive data: Data) {
        buffer.append(data)
        
        // 尝试将缓冲区分割成行
        if let string = String(data: buffer, encoding: .utf8) {
            let lines = string.components(separatedBy: "\n")
            
            // 处理除最后一行外的所有行（最后一行可能不完整）
            for i in 0..<lines.count-1 {
                let line = lines[i]
                processLine(line)
            }
            
            // 保留最后一行，它可能不完整
            if let lastLine = lines.last, !lastLine.isEmpty {
                buffer = lastLine.data(using: .utf8) ?? Data()
            } else {
                buffer = Data()
            }
        }
    }
    
    func urlSession(_ session: URLSession, task: URLSessionTask, didCompleteWithError error: Error?) {
        // 处理剩余的缓冲区
        if let string = String(data: buffer, encoding: .utf8), !string.isEmpty {
            processLine(string)
        }
        
        onComplete(error)
        self.task = nil
    }
    
    private func processLine(_ line: String) {
        if line.hasPrefix("data: ") {
            let jsonString = line.dropFirst(6) // 去掉 "data: "
            
            if jsonString == "[DONE]" {
                return
            }
            
            do {
                if let data = String(jsonString).data(using: .utf8),
                   let json = try JSONSerialization.jsonObject(with: data) as? [String: Any],
                   let choices = json["choices"] as? [[String: Any]],
                   let choice = choices.first,
                   let delta = choice["delta"] as? [String: Any],
                   let content = delta["content"] as? String {
                    
                    onReceive(content)
                }
            } catch {
                print("解析JSON错误: \(error)")
            }
        }
    }
} 
